import numpy

# -----------------------------
# ABSTRACT CLASSES
# -----------------------------

class Module:
	def __init__(self): pass
	def update(self,lr): pass

class Criterion:
	def __init__(self): pass

# -----------------------------
# CONTAINERS MODULES
# -----------------------------

class Sequential(Module):

	def __init__(self,modules):
		self.modules = modules

	def forward(self,X):
		for m in self.modules:
			X = m.forward(X)
		return X

	def backward(self,DY):
		for m in self.modules[::-1]:
			DY = m.backward(DY)
		return DY

	def update(self,lr):
		for m in self.modules:
			m.update(lr)

# -----------------------------
# LEARNABLE MODULES
# -----------------------------

class Linear(Module):

	def __init__(self,m,n,zero=False):

		self.tr = m**.5 / n**.5
		self.lr = 1 / m**.5
		
		self.W = numpy.random.normal(0,1 / m**.5,[m,n]).astype('float32')
		self.B = numpy.zeros([n]).astype('float32')

		if zero: self.W *= 0

	def forward(self,X):
		self.X = X
		return numpy.dot(X,self.W)+self.B

	def backward(self,DY):
		self.DW = numpy.dot(self.X.T,DY)
		self.DB = DY.sum(axis=0)
		return self.tr*numpy.dot(DY,self.W.T)

	def update(self,lr):
		self.W -= lr*self.lr*self.DW
		self.B -= lr*self.lr*self.DB


# -----------------------------
# TRANSFER FUNCTIONS
# -----------------------------

class Sigmoid(Module): # Sigmoid transfer function: f(x) = 1.5*tanh(x)

	def forward(self,X):   pass # TODO: implement this function
	def backward(self,DY): pass # TODO: implement this function


class Positive(Module): # Positive-part transfer function: f(x) = max(0,x)

	def forward(self,X):
		self.X = X
		return numpy.maximum(self.X,0)

	def backward(self,DY):
		return DY * (self.X > 0)

# -----------------------------
# CRITERIONS
# -----------------------------

class SquareError(Criterion): # Squared error criterion
	def grad(self,Y,T): return Y-T

class NegLogLik(Criterion): # Negative log-likelihood criterion
	def grad(self,Y,T): return numpy.exp(Y)/numpy.exp(Y).sum(axis=1)[:,numpy.newaxis]-T

